from __future__ import print_function, division
import shutil
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from torchvision import models
from lib.network_xh import Network
import torchvision
import vis_nlp
import numpy as np
import cv2
import math
import os


def vis_nl_map(img_path, nl_map, bt_idx, vis_size=(560, 560)):
    img = cv2.imread(img_path, 1)
    img = cv2.resize(img, dsize=vis_size)
    h, w, c = img.shape
    nl_map_0 = nl_map[0]
    total_region, nl_map_length = nl_map_0.shape
    region_per_row = round(math.sqrt(total_region))
    size_of_region = round(w / region_per_row)

    nl_map_size = round(math.sqrt(nl_map_length))

    for index in range(total_region):
        img_draw = img.copy()
        nl_map = nl_map_0[index]
        nl_map = nl_map.reshape(nl_map_size, nl_map_size)
        nl_map = cv2.resize(nl_map, dsize=(h, w))
        nl_map = np.uint8(nl_map * 255)
        heat_img = cv2.applyColorMap(nl_map, cv2.COLORMAP_JET)
        heat_img = cv2.cvtColor(heat_img, cv2.COLOR_BGR2RGB)
        img_add = cv2.addWeighted(img_draw, 0.3, heat_img, 0.7, 0)
        x0 = index // region_per_row * size_of_region
        x1 = x0 + size_of_region
        y0 = index % region_per_row * size_of_region
        y1 = y0 + size_of_region
        cv2.rectangle(img_add, (y0, x0), (y1, x1), (255, 0, 0), 1)
        cv2.imshow("nlp", img_add)
        cv2.imwrite("./nlp/%d_%d.jpg"%(bt_idx, index), img_add)
        cv2.waitKey(10)


def del_dir(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)       # 取文件绝对路径
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_dir(path_file)


def rm_mkdir(dir_path):
    if os.path.exists(dir_path):
        del_dir(dir_path)
        print('Clean path - %s' % dir_path)
    else:
        os.makedirs(dir_path)
        print('Create path - %s' % dir_path)


def load_model(test_dir):
    # test_datasets = torchvision.datasets.MNIST(root='./mnist/',
    #                                        transform=torchvision.transforms.ToTensor(),
    #                                        train=False)
    val_transforms = transforms.Compose([
        transforms.Grayscale(),
        transforms.ToTensor()
    ])
    test_datasets = datasets.ImageFolder(test_dir, transform=val_transforms)

    net = Network()
    if torch.cuda.is_available():
        net = nn.DataParallel(net)
        net.cuda()
    state_dict = torch.load(r"./weights/xh_net.pth")
    net.load_state_dict(state_dict)
    return net, test_datasets


def run():
    test_dir = r'E:\xianhuang0818\10_cls_ori\for_train\test'
    net, test_datasets = load_model(test_dir)
    test_dataloader = torch.utils.data.DataLoader(test_datasets, batch_size=1, shuffle=False)
    rm_mkdir("./result")
    name_list_en = os.listdir(test_dir)
    bt_idx = 0
    with torch.no_grad():
        correct = 0
        total = 0
        err_idx = 0
        cls_idx = test_datasets.class_to_idx
        cls = test_datasets.classes
        class_num = len(cls)
        conf_mat = np.zeros((class_num, class_num), np.uint32)
        for data in test_dataloader:
            net.eval()
            inputs_ori, labels_ori = data
            inputs, labels = Variable(inputs_ori).cuda(), Variable(labels_ori).cuda()
            _, nl_mep_list = net.module.forward_with_nl_map(inputs)
            # (b, h1*w1, h2*w2)
            nl_map_1 = nl_mep_list[0].cpu().numpy()
            nl_map_2 = nl_mep_list[1].cpu().numpy()
            nl_map_3 = nl_mep_list[2].cpu().numpy()
            nl_map_4 = nl_mep_list[3].cpu().numpy()
            nl_map_5 = nl_mep_list[4].cpu().numpy()

            img = torchvision.transforms.ToPILImage()(inputs.cpu()[0])
            img.save('nl_map_vis/sample.png')
            # np.save('nl_map_vis/nl_map_1', nl_map_1)
            # np.save('nl_map_vis/nl_map_2', nl_map_2)
            # vis_nl_map(img_path='nl_map_vis/sample.png', nl_map=nl_map_1, bt_idx=bt_idx, vis_size=(560, 560))
            vis_nl_map(img_path='nl_map_vis/sample.png', nl_map=nl_map_1, bt_idx=bt_idx, vis_size=(560, 560))

            bt_idx += 1
            # np.save('nl_map_vis/nl_map_1_%d'%bt_idx, nl_map_1)
            # np.save('nl_map_vis/nl_map_2_%d'%bt_idx, nl_map_2)


if __name__ == "__main__":
    run()
    # train()